package main

// reference https://books.halfrost.com/leetcode/ChapterThree/LFUCache/#%e8%a7%a3%e6%b3%95%e4%ba%8c-get-ocapacity--put-ocapacity

import "container/heap"

type (
	Item struct {
		value, key, frequency, count, index int
	}
	PriorityQueue []*Item

	LFUCache struct {
		capacity int
		pq       PriorityQueue
		hash     map[int]*Item
		counter  int
	}
)

func (pq PriorityQueue) Len() int {
	return len(pq)
}

func (pq PriorityQueue) Less(i, j int) bool {
	if pq[i].frequency == pq[j].frequency {
		return pq[i].count < pq[j].count
	}
	return pq[i].frequency < pq[j].frequency
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *PriorityQueue) Push(x any) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() any {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil
	item.index = -1
	*pq = old[0 : n-1]
	return item
}

func (pq *PriorityQueue) update(item *Item, value, frequency, count int) {
	item.value = value
	item.count = count
	item.frequency = frequency
	heap.Fix(pq, item.index)
}

func (c *LFUCache) Get(key int) int {
	if c.capacity == 0 {
		return -1
	}
	if item, ok := c.hash[key]; ok {
		c.counter++
		c.pq.update(item, item.value, item.frequency+1, c.counter)
		return item.value
	}
	return -1
}

func (c *LFUCache) Put(key, value int) {
	if c.capacity == 0 {
		return
	}
	c.counter++
	if item, ok := c.hash[key]; ok {
		c.pq.update(item, value, item.frequency+1, c.counter)
		return
	}

	if len(c.pq) == c.capacity {
		item := heap.Pop(&c.pq).(*Item)
		delete(c.hash, item.key)
	}

	item := &Item{
		value: value,
		key:   key,
		count: c.counter,
	}
	heap.Push(&c.pq, item)
	c.hash[key] = item
}

func Constructor(capacity int) LFUCache {
	return LFUCache{
		pq:       PriorityQueue{},
		hash:     make(map[int]*Item, capacity),
		capacity: capacity,
	}
}
